package drpc.complex;

import org.apache.storm.Config;
import org.apache.storm.LocalCluster;
import org.apache.storm.LocalDRPC;
import org.apache.storm.StormSubmitter;
import org.apache.storm.coordination.BatchOutputCollector;
import org.apache.storm.drpc.LinearDRPCTopologyBuilder;
import org.apache.storm.task.TopologyContext;
import org.apache.storm.topology.BasicOutputCollector;
import org.apache.storm.topology.OutputFieldsDeclarer;
import org.apache.storm.topology.base.BaseBasicBolt;
import org.apache.storm.topology.base.BaseBatchBolt;
import org.apache.storm.tuple.Fields;
import org.apache.storm.tuple.Tuple;
import org.apache.storm.tuple.Values;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.*;

public class ReachTopology {
    private static final Logger logger = LoggerFactory.getLogger(ReachTopology.class);

    private static final Map<String, List<String>> TWEETERS_DB = new HashMap<String, List<String>>(){
        {
            put("foo.com/blog/1", Arrays.asList("sally", "bob", "tim", "george", "nathan"));
            put("engineering.twitter.com/blog/5", Arrays.asList("adam", "david", "sally", "nathan"));
            put("tech.backtype.com/blog/123", Arrays.asList("tim", "mike", "john"));
        }
    };

    private static final Map<String, List<String>> FOLLOWERS_DB = new HashMap<String, List<String>>(){
        {
            put("sally",Arrays.asList("bob", "tim", "alice", "adam", "jim", "chris", "jai"));
            put("bob",Arrays.asList("sally", "nathan", "jim", "mary", "david", "vivian"));
            put("tim", Collections.singletonList("alex"));
            put("nathan",Arrays.asList("sally", "bob", "adam", "harry", "chris", "vivian", "emily", "jordan"));
            put("adam",Arrays.asList("david", "carissa"));
            put("mike",Arrays.asList("john", "bob"));
            put("john",Arrays.asList("alice", "nathan", "jim", "mike", "bob"));
        }
    };

    public static class GetTweeters extends BaseBasicBolt {

        @Override
        public void execute(Tuple tuple, BasicOutputCollector basicOutputCollector) {

            List<String> strings = TWEETERS_DB.get(tuple.getString(1));

            if (strings == null) {
                return;
            }

            for (String user : strings) {
                basicOutputCollector.emit(new Values(tuple.getValue(0), user));
            }
        }

        @Override
        public void declareOutputFields(OutputFieldsDeclarer outputFieldsDeclarer) {
            outputFieldsDeclarer.declare(new Fields("id", "tweeter"));
        }
    }


    private static class GetFollower extends BaseBasicBolt {
        @Override
        public void execute(Tuple tuple, BasicOutputCollector basicOutputCollector) {
            String tweeter = tuple.getString(1);
            List<String> followers = FOLLOWERS_DB.get(tweeter);
            if (followers == null) {
                return;
            }

            for (String follower: followers) {
                basicOutputCollector.emit(new Values(tuple.getValue(0), follower));
            }
        }

        @Override
        public void declareOutputFields(OutputFieldsDeclarer outputFieldsDeclarer) {
            outputFieldsDeclarer.declare(new Fields("id", "follower"));
        }
    }


    /**
     * 提供将一批Tuple作为一个整体进行处理的API
     * 每一个request id都会创建一个batch bolt实例，Storm会在合适的时候负责清理该实例
     */
    public static class PartialUniquer extends BaseBatchBolt<Object> {
        private BatchOutputCollector collector;
        private Set<String> followers;
        private Object id;

        @Override
        public void prepare(Map map, TopologyContext topologyContext, BatchOutputCollector batchOutputCollector, Object id) {
            collector = batchOutputCollector;
            followers = new HashSet<>();
            this.id = id;
        }

        @Override
        public void execute(Tuple tuple) {
            followers.add(tuple.getString(1));
        }

        @Override
        public void finishBatch() {
            collector.emit(new Values(id, followers.size()));
        }

        @Override
        public void declareOutputFields(OutputFieldsDeclarer outputFieldsDeclarer) {
            outputFieldsDeclarer.declare(new Fields("id", "partial-count"));
        }
    }

    public static class CountAggregator extends BaseBatchBolt<Object> {
        private BatchOutputCollector collector;
        private Object id;
        private int count;

        @Override
        public void prepare(Map<String, Object> map, TopologyContext topologyContext, BatchOutputCollector batchOutputCollector, Object o) {
            collector = batchOutputCollector;
            this.id = o;
            count = 0;
        }

        @Override
        public void execute(Tuple tuple) {
            count += tuple.getInteger(1);
        }

        @Override
        public void finishBatch() {
            collector.emit(new Values(id, count));
        }

        @Override
        public void declareOutputFields(OutputFieldsDeclarer outputFieldsDeclarer) {
            outputFieldsDeclarer.declare(new Fields("id", "reach"));
        }
    }

    public static void main(String[] args) throws Exception {
        String function = "reach";
        String topology = "reach-drpc";

        // 此处定义函数名
        LinearDRPCTopologyBuilder builder = new LinearDRPCTopologyBuilder(function);

        builder.addBolt(new GetTweeters(), 2).shuffleGrouping();
        builder.addBolt(new GetFollower(), 4).shuffleGrouping();
        builder.addBolt(new PartialUniquer(), 6).fieldsGrouping(new Fields("id", "follower"));
        builder.addBolt(new CountAggregator()).fieldsGrouping(new Fields("id"));
        // TODO 再写一个bolt，将每个分区的加起来

        Config config = new Config();
        if (args == null || args.length == 0) {
            LocalDRPC localDRPC = new LocalDRPC();
            LocalCluster localCluster = new LocalCluster();

            localCluster.submitTopology(topology, config, builder.createLocalTopology(localDRPC));

            String[] urls = new String[]{
                    "foo.com/blog/1", "engineering.twitter.com/blog/5", "baidu.com"
            };

            for (String url : urls) {
                String result = localDRPC.execute(function, url);
                logger.info("---------------------------------------" + url + " -- " + result);
            }
            localCluster.shutdown();
            localDRPC.shutdown();
        } else {
            topology = args[0];
            StormSubmitter.submitTopology(topology, config, builder.createRemoteTopology());
        }
    }
}
